from django.db import models
from django.contrib.auth import get_user_model
from django.conf import settings
from apps.common.models import BaseModel
import base64

User = get_user_model()

# 导入流式会话模型
from .models_streaming import StreamSession



class AIProvider(BaseModel):
    """AI服务提供商配置表"""

    PROVIDER_TYPES = [
        ('openai', 'OpenAI'),
        ('anthropic', 'Anthropic'),
        ('google', 'Google'),
        ('custom', '自定义'),
    ]

    # 业务用户关联字段 - 使用不同的related_name
    owner = models.ForeignKey(
        User,
        on_delete=models.CASCADE,
        null=True,
        blank=True,
        related_name='owned_ai_providers',
        verbose_name="所有者",
        help_text="NULL表示系统预设提供商"
    )
    name = models.CharField(max_length=100, verbose_name="提供商名称")
    provider_type = models.CharField(
        max_length=20,
        choices=PROVIDER_TYPES,
        verbose_name="提供商类型"
    )
    api_endpoint = models.URLField(verbose_name="API端点")
    api_key = models.TextField(verbose_name="API密钥(加密存储)")
    default_params = models.JSONField(
        default=dict,
        blank=True,
        verbose_name="默认参数JSON"
    )
    is_active = models.BooleanField(default=True, verbose_name="是否启用")
    is_public = models.BooleanField(
        default=False,
        verbose_name="是否公开",
        help_text="系统预设提供商为公开"
    )

    class Meta:
        verbose_name = "AI提供商"
        verbose_name_plural = "AI提供商"
        unique_together = ['owner', 'name']
        ordering = ['-created_at']
        indexes = [
            models.Index(fields=['owner', 'is_active']),
            models.Index(fields=['provider_type', 'is_active']),
        ]

    def __str__(self):
        return f"{self.name} ({self.get_provider_type_display()})"

    def encrypt_api_key(self, api_key):
        """加密API密钥"""
        # 使用简单的base64编码作为示例，生产环境应使用更安全的加密方式
        return base64.b64encode(api_key.encode()).decode()

    def decrypt_api_key(self):
        """解密API密钥"""
        try:
            return base64.b64decode(self.api_key.encode()).decode()
        except:
            return self.api_key  # 如果解密失败，返回原值

    def save(self, *args, **kwargs):
        # 如果是新记录或API密钥发生变化，进行加密
        if not self.pk or 'api_key' in kwargs.get('update_fields', []):
            # 简单检查是否已加密（base64编码后的字符串特征）
            try:
                base64.b64decode(self.api_key)
            except:
                # 未加密，进行加密
                self.api_key = self.encrypt_api_key(self.api_key)
        super().save(*args, **kwargs)

    def can_edit(self, user):
        """权限检查：是否可以编辑"""
        return self.owner == user or user.is_staff

    def can_delete(self, user):
        """权限检查：是否可以删除"""
        return self.owner == user or user.is_staff


class AIModel(BaseModel):
    """AI模型配置表"""

    provider = models.ForeignKey(
        AIProvider,
        on_delete=models.CASCADE,
        related_name='models',
        verbose_name="提供商"
    )
    name = models.CharField(max_length=100, verbose_name="模型显示名称")
    model_code = models.CharField(
        max_length=100,
        verbose_name="模型代码",
        help_text="如gpt-4, claude-3-opus"
    )
    max_tokens = models.IntegerField(
        default=4096,
        verbose_name="最大Token限制"
    )
    cost_per_1k_tokens = models.DecimalField(
        max_digits=10,
        decimal_places=6,
        default=0.000000,
        verbose_name="每1k Token成本(参考)"
    )
    model_params = models.JSONField(
        default=dict,
        blank=True,
        verbose_name="模型特定参数JSON"
    )
    is_active = models.BooleanField(default=True, verbose_name="是否启用")

    class Meta:
        verbose_name = "AI模型"
        verbose_name_plural = "AI模型"
        unique_together = ['provider', 'model_code']
        ordering = ['-created_at']
        indexes = [
            models.Index(fields=['provider', 'is_active']),
            models.Index(fields=['model_code']),
        ]

    def __str__(self):
        return f"{self.name} ({self.model_code})"

    def can_edit(self, user):
        """权限检查：是否可以编辑"""
        return self.provider.can_edit(user)

    def can_delete(self, user):
        """权限检查：是否可以删除"""
        return self.provider.can_delete(user)


class AIRole(BaseModel):
    """AI角色配置表"""

    CATEGORIES = [
        ('assistant', '通用助手'),
        ('creative', '创作写作'),
        ('analysis', '分析总结'),
        ('education', '教育培训'),
        ('programming', '编程开发'),
        ('business', '商务办公'),
        ('entertainment', '娱乐休闲'),
        ('custom', '自定义'),
    ]

    name = models.CharField(max_length=100, verbose_name="角色名称")
    system_prompt = models.TextField(verbose_name="系统提示词")
    category = models.CharField(
        max_length=20,
        choices=CATEGORIES,
        default='custom',
        verbose_name="角色分类"
    )
    is_public = models.BooleanField(
        default=False,
        verbose_name="是否公开"
    )
    usage_count = models.IntegerField(
        default=0,
        verbose_name="使用次数"
    )

    class Meta:
        verbose_name = "AI角色"
        verbose_name_plural = "AI角色"
        unique_together = ['created_by', 'name']
        ordering = ['-created_at']
        indexes = [
            models.Index(fields=['created_by', 'is_public']),
            models.Index(fields=['category']),
            models.Index(fields=['-usage_count']),
        ]

    def __str__(self):
        return f"{self.name} ({self.get_category_display()})"

    def can_edit(self, user):
        """权限检查：是否可以编辑"""
        return self.created_by == user or user.is_staff

    def can_delete(self, user):
        """权限检查：是否可以删除"""
        return self.created_by == user or user.is_staff


class AIConversation(BaseModel):
    """AI对话会话表"""

    STATUS_CHOICES = [
        ('active', '活跃'),
        ('archived', '已归档'),
        ('deleted', '已删除'),
    ]

    # 业务用户关联字段 - 使用不同的related_name
    owner = models.ForeignKey(
        User,
        on_delete=models.CASCADE,
        related_name='owned_conversations',
        verbose_name="所有者"
    )
    ai_model = models.ForeignKey(
        AIModel,
        on_delete=models.SET_NULL,
        null=True,
        related_name='conversations',
        verbose_name="AI模型"
    )
    ai_role = models.ForeignKey(
        AIRole,
        on_delete=models.SET_NULL,
        null=True,
        related_name='conversations',
        verbose_name="AI角色"
    )
    title = models.CharField(max_length=200, verbose_name="会话标题")
    used_tokens = models.IntegerField(
        default=0,
        verbose_name="已使用Token数"
    )
    status = models.CharField(
        max_length=20,
        choices=STATUS_CHOICES,
        default='active',
        verbose_name="状态"
    )
    last_message_at = models.DateTimeField(
        null=True,
        blank=True,
        verbose_name="最后消息时间"
    )

    class Meta:
        verbose_name = "AI对话会话"
        verbose_name_plural = "AI对话会话"
        ordering = ['-last_message_at', '-created_at']
        indexes = [
            models.Index(fields=['owner', 'status']),
            models.Index(fields=['-last_message_at']),
            models.Index(fields=['ai_model']),
            models.Index(fields=['ai_role']),
        ]

    def __str__(self):
        return f"{self.title} - {self.owner.username}"

    def can_edit(self, user):
        """权限检查：是否可以编辑"""
        return self.owner == user or user.is_staff

    def can_delete(self, user):
        """权限检查：是否可以删除"""
        return self.owner == user or user.is_staff


class AIMessage(BaseModel):
    """对话消息表"""

    conversation = models.ForeignKey(
        AIConversation,
        on_delete=models.CASCADE,
        related_name='messages',
        verbose_name="会话"
    )
    ai_role = models.ForeignKey(
        AIRole,
        on_delete=models.SET_NULL,
        null=True,
        blank=True,
        related_name='messages',
        verbose_name="AI角色",
        help_text="NULL表示用户消息"
    )
    content = models.TextField(verbose_name="消息内容")
    token_count = models.IntegerField(
        default=0,
        verbose_name="Token数量"
    )
    metadata = models.JSONField(
        default=dict,
        blank=True,
        verbose_name="元数据",
        help_text="模型参数等"
    )
    parent_message = models.ForeignKey(
        'self',
        on_delete=models.SET_NULL,
        null=True,
        blank=True,
        related_name='child_messages',
        verbose_name="父消息ID",
        help_text="用于分支对话"
    )

    class Meta:
        verbose_name = "对话消息"
        verbose_name_plural = "对话消息"
        ordering = ['created_at']
        indexes = [
            models.Index(fields=['conversation', 'created_at']),
            models.Index(fields=['ai_role']),
            models.Index(fields=['parent_message']),
        ]

    def __str__(self):
        role = "AI" if self.ai_role else "用户"
        return f"{role}: {self.content[:50]}..."

    @property
    def is_user_message(self):
        """是否为用户消息"""
        return self.ai_role is None

    def can_edit(self, user):
        """权限检查：是否可以编辑"""
        return self.conversation.can_edit(user)

    def can_delete(self, user):
        """权限检查：是否可以删除"""
        return self.conversation.can_delete(user)


class TokenUsageLog(BaseModel):
    """用户Token使用统计表"""

    # 业务用户关联字段 - 使用不同的related_name
    user = models.ForeignKey(
        User,
        on_delete=models.CASCADE,
        related_name='token_usage_logs',
        verbose_name="用户"
    )
    conversation = models.ForeignKey(
        AIConversation,
        on_delete=models.CASCADE,
        related_name='token_logs',
        verbose_name="会话"
    )
    message = models.ForeignKey(
        AIMessage,
        on_delete=models.CASCADE,
        related_name='token_logs',
        verbose_name="消息"
    )
    ai_model = models.ForeignKey(
        AIModel,
        on_delete=models.CASCADE,
        related_name='token_logs',
        verbose_name="AI模型"
    )
    input_tokens = models.IntegerField(
        default=0,
        verbose_name="输入Token数"
    )
    output_tokens = models.IntegerField(
        default=0,
        verbose_name="输出Token数"
    )
    total_tokens = models.IntegerField(
        default=0,
        verbose_name="总Token数"
    )

    class Meta:
        verbose_name = "Token使用记录"
        verbose_name_plural = "Token使用记录"
        ordering = ['-created_at']
        indexes = [
            models.Index(fields=['user', 'created_at']),
            models.Index(fields=['ai_model', 'created_at']),
            models.Index(fields=['conversation']),
        ]

    def __str__(self):
        return f"{self.user.username} - {self.total_tokens} tokens"

    def save(self, *args, **kwargs):
        # 自动计算总Token数
        self.total_tokens = self.input_tokens + self.output_tokens
        super().save(*args, **kwargs)


class MessageTemplate(BaseModel):
    """消息模板表"""

    # 业务用户关联字段 - 使用不同的related_name
    owner = models.ForeignKey(
        User,
        on_delete=models.CASCADE,
        related_name='message_templates',
        verbose_name="所有者"
    )
    name = models.CharField(max_length=100, verbose_name="模板名称")
    content = models.TextField(verbose_name="模板内容")
    variables = models.JSONField(
        default=dict,
        blank=True,
        verbose_name="变量定义",
        help_text="JSON格式的变量定义"
    )
    category = models.CharField(
        max_length=50,
        default='general',
        verbose_name="分类"
    )
    is_public = models.BooleanField(
        default=False,
        verbose_name="是否公开"
    )
    usage_count = models.IntegerField(
        default=0,
        verbose_name="使用次数"
    )

    class Meta:
        verbose_name = "消息模板"
        verbose_name_plural = "消息模板"
        unique_together = ['owner', 'name']
        ordering = ['-usage_count', '-created_at']
        indexes = [
            models.Index(fields=['owner', 'category']),
            models.Index(fields=['is_public']),
            models.Index(fields=['-usage_count']),
        ]

    def __str__(self):
        return f"{self.name} - {self.owner.username}"

    def can_edit(self, user):
        """权限检查：是否可以编辑"""
        return self.owner == user or user.is_staff

    def can_delete(self, user):
        """权限检查：是否可以删除"""
        return self.owner == user or user.is_staff


class UserFavorite(BaseModel):
    """用户收藏表"""

    # 业务用户关联字段 - 使用不同的related_name
    owner = models.ForeignKey(
        User,
        on_delete=models.CASCADE,
        related_name='favorites',
        verbose_name="所有者"
    )
    conversation = models.ForeignKey(
        AIConversation,
        on_delete=models.CASCADE,
        null=True,
        blank=True,
        related_name='favorites',
        verbose_name="会话"
    )
    message = models.ForeignKey(
        AIMessage,
        on_delete=models.CASCADE,
        null=True,
        blank=True,
        related_name='favorites',
        verbose_name="消息"
    )
    title = models.CharField(max_length=200, verbose_name="收藏标题")
    note = models.TextField(
        blank=True,
        verbose_name="备注"
    )

    class Meta:
        verbose_name = "用户收藏"
        verbose_name_plural = "用户收藏"
        unique_together = [
            ['owner', 'conversation'],
            ['owner', 'message'],
        ]
        ordering = ['-created_at']
        indexes = [
            models.Index(fields=['owner', 'created_at']),
            models.Index(fields=['conversation']),
            models.Index(fields=['message']),
        ]

    def __str__(self):
        return f"{self.title} - {self.owner.username}"

    def can_edit(self, user):
        """权限检查：是否可以编辑"""
        return self.owner == user or user.is_staff

    def can_delete(self, user):
        """权限检查：是否可以删除"""
        return self.owner == user or user.is_staff
